package my.algo.data;

/**
 * @author fengbo
 * @date 2018/8/7
 */
public class TreeMap<K extends Comparable<K>, V> implements Map<K, V> {

    private static final boolean RED = true;
    private static final boolean BLACK = false;

    private Node<K, V> root;

    @Override
    public void put(K key, V value) {
        root = put(root, key, value);
        root.color = BLACK;
    }

    private Node<K, V> put(Node<K, V> node, K key, V value) {
        if (node == null) {
            return new Node<>(key, value, 1, RED);
        }

        int cmp = key.compareTo(node.key);
        if (cmp < 0) {
            node.left = put(node.left, key, value);
        } else if (cmp > 0) {
            node.right = put(node.right, key, value);
        } else {
            node.value = value;
        }

        if (isRed(node.right) && !isRed(node.left)) {
            node = rotateLeft(node);
        }
        if (isRed(node.left) && isRed(node.left.left)) {
            node = rotateRight(node);
        }
        if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }
        node.n = size(node.left) + size(node.right) + 1;
        return node;
    }

    private boolean isRed(Node<K, V> node) {
        return node != null && node.color == RED;
    }

    private Node<K, V> rotateLeft(Node<K, V> node) {
        Node<K, V> x = node.right;
        node.right = x.left;
        x.left = node;
        x.color = node.color;
        node.color = RED;
        x.n = node.n;
        node.n = size(node.right) + size(node.left) + 1;
        return x;
    }

    private Node<K, V> rotateRight(Node<K, V> node) {
        Node<K, V> x = node.left;
        node.left = x.right;
        x.right = node;
        x.color = node.color;
        node.color = RED;
        x.n = node.n;
        node.n = size(node.left) + size(node.right) + 1;
        return x;
    }

    private void flipColors(Node<K, V> node) {
        node.color = RED;
        node.left.color = BLACK;
        node.right.color = BLACK;
    }

    @Override
    public V get(K key) {
        return get(root, key);
    }

    private V get(Node<K, V> node, K key) {
        if (node == null) {
            return null;
        }

        int cmp = key.compareTo(node.key);
        if (cmp > 0) {
            return get(node.right, key);
        } else if (cmp < 0) {
            return get(node.left, key);
        } else {
            return node.value;
        }
    }

    @Override
    public void delete(K key) {
        if (!isRed(root.left) && !isRed(root.right)) {
            root.color = RED;
        }
        root = delete(root, key);
        if (!isEmpty()) {
            root.color = BLACK;
        }
    }

    private Node<K, V> delete(Node<K, V> node, K key) {
        if (key.compareTo(node.key) < 0) {
            if (!isRed(node.left) && !isRed(node.left.left)) {
                node = moveRedLeft(node);
            }
            node.left = delete(node.left, key);
        } else {
            if (isRed(node.left)) {
                node = rotateRight(node);
            }
            if (key.compareTo(node.key) == 0 && (node.right == null)) {
                return null;
            }
            if (!isRed(node.right) && !isRed(node.right.left)) {
                node = moveRedRight(node);
            }
            if (key.compareTo(node.key) == 0) {
                node.value = get(node.right, min(node.right).key);
                node.key = min(node.right).key;
                node.right = resizeRight(node.right);
            } else {
                node.right = delete(node.right, key);
            }
        }
        return balance(node);
    }

    private Node<K, V> balance(Node<K, V> node) {
        if (isRed(node.right)) {
            node = rotateLeft(node);
        }
        if (isRed(node.right) && !isRed(node.left)) {
            node = rotateLeft(node);
        }
        if (isRed(node.left) && isRed(node.left.left)) {
            node = rotateRight(node);
        }
        if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }
        node.n = size(node.left) + size(node.right) + 1;
        return node;
    }

    private Node<K, V> moveRedLeft(Node<K, V> node) {
        flipColors(node);
        if (isRed(node.right.left)) {
            node.right = rotateRight(node.right);
            node = rotateLeft(node);
        }
        return node;
    }

    private Node<K, V> moveRedRight(Node<K, V> node) {
        flipColors(node);
        if (!isRed(node.left.left)) {
            node = rotateRight(node);
        }
        return node;
    }

    private Node<K, V> resizeRight(Node<K, V> node) {
        if (node.left == null) {
            return node.right;
        }
        node.left = resizeRight(node.left);
        node.n = size(node.left) + size(node.right) + 1;
        return node;
    }

    private Node<K, V> min(Node<K, V> node) {
        if (node.left == null) {
            return node;
        }
        return min(node.left);
    }

    @Override
    public boolean contains(K key) {
        return get(key) != null;
    }

    @Override
    public boolean isEmpty() {
        return root != null && root.n > 0;
    }

    @Override
    public int size() {
        return size(root);
    }

    private int size(Node node) {
        if (node == null) {
            return 0;
        }
        return node.n;
    }

    @Override
    public List<K> keys() {
        List<K> result = new LinkedList<>();
        return keys(root, result);
    }

    private List<K> keys(Node<K, V> node, List<K> keys) {
        if (node == null) {
            return keys;
        }
        keys(node.left, keys);
        keys.add(node.key);
        keys(node.right, keys);
        return keys;
    }

    private static class Node<K extends Comparable<K>, V> {
        K key;
        V value;
        int n;
        boolean color;
        Node<K, V> left, right;
        Node (K key, V value, int n, boolean color) {
            this.key = key;
            this.value = value;
            this.n = n;
            this.color = color;
        }
    }

    public static void main(String[] args) {
        Map<Integer, Integer> map = new TreeMap<>();
        map.put(10, 10);
        map.put(10, 10);
        map.put(20, 20);
        map.put(30, 30);
        map.put(40, 40);
        map.put(50, 50);
        map.put(60, 60);
        map.put(70, 70);
        map.put(80, 80);
        map.put(90, 90);
        map.put(100, 100);
        map.put(110, 110);
        map.put(120, 120);
        map.put(130, 130);

        List<Integer> keys = map.keys();
        for (Integer key : keys) {
            System.out.println(map.get(key));
        }
    }
}
